from modelscope import AutoTokenizer, AutoModelForCausalLM
import streamlit as st
from transformers import TextStreamer, TextIteratorStreamer
from threading import Thread



# 在侧边栏创建标题
with st.sidebar:
    st.markdown("qwen2")
    # 创建滑块，默认值为512，范围在0到1024之间
    max_length = st.slider("max_length", 0, 1024, 512, step=1)
# 创建标题和副标题
# st.title("qwen2 chatbot")
# st.caption("test")
# 你下载到本地的模型路径

model_name = 'Qwen/Qwen1.5-0.5B-Chat'
# model_name = 'Qwen/Qwen2.5-Coder-1.5B-Instruct-AWQ'


# @streamlit.cache_resource 是一个用于缓存昂贵或频繁调用的资源（如大型文件、网络资源、或数据库连接）的装饰器。这个装饰器可以帮助你提高应用的性能，通过缓存那些不经常变更但加载需要大量时间或计算资源的数据。
# 定义的函数来获取tokenizer和model
@st.cache_resource
def get_model():
    tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype="auto",
        device_map="auto"
    )
    # model.to(torch.device('cpu:1,7') )
    return tokenizer, model


tokenizer, model = get_model()
# 如果没有消息，则创建默认的消息列表
if "messages" not in st.session_state:
    st.session_state["messages"] = [
        {"role": "assistant", "content": "有什么可以帮到您？"}
    ]
# 便利session_state中的消息并显示在聊天界面上
for msg in st.session_state.messages:
    st.chat_message(msg["role"]).write(msg["content"])
## 如果用户在聊天输入框中输入了内容，则执行下述操作
if prompt := st.chat_input():
    # 将用户输入添加到message列表中
    st.session_state.messages.append({"role": "user", "content": prompt})
    # 在聊天界面上显示用户输入
    st.chat_message("user").write(prompt)

    # 构建输入
    input_ids = tokenizer.apply_chat_template(
        st.session_state.messages, tokenize=False, add_generation_prompt=True
    )
    model_inputs = tokenizer([input_ids], return_tensors="pt").to(model.device)
    # 模型生成输出id
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    generation_kwargs = dict(model_inputs, streamer=streamer, max_new_tokens=512)
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    response = st.chat_message("assistant").write_stream(streamer)
    st.session_state.messages.append({"role": "assistant", "content": response})
    # 在界面上显示输出
